package mrpanyu.guitool.jwt;

import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Date;
import java.util.Map;

import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import com.fasterxml.jackson.databind.ObjectMapper;

import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jws;
import io.jsonwebtoken.JwtBuilder;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.security.MacAlgorithm;
import io.jsonwebtoken.security.SecureDigestAlgorithm;
import mrpanyu.guitool.base.annotation.Action;
import mrpanyu.guitool.base.annotation.OnParameterChange;
import mrpanyu.guitool.base.annotation.Parameter;
import mrpanyu.guitool.base.annotation.ToolModel;
import mrpanyu.guitool.base.model.ParameterType;
import mrpanyu.guitool.base.model.Tool;
import mrpanyu.guitool.base.util.CommonUtils;

@ToolModel(displayName = "Jwt生成/解析")
public class JwtTool {

	@Parameter(displayName = "算法", type = ParameterType.SELECT, options = { "HS256", "HS384", "HS512", "RS256", "RS384",
			"RS512", "ES256", "ES384", "ES512" }, order = 1)
	private String algorithm = "HS256";

	@Parameter(displayName = "操作", type = ParameterType.SELECT, options = { "Generate", "Parse" }, order = 2)
	private String operation = "Generate";

	@Parameter(displayName = "密钥", description = "Base64编码", type = ParameterType.TEXT, order = 3)
	private String key;

	@Parameter(displayName = "私钥", description = "PKCS8格式Base64", type = ParameterType.TEXT, order = 4, visible = false)
	private String privateKey;

	@Parameter(displayName = "公钥", description = "X509格式Base64", type = ParameterType.TEXT, order = 5, visible = false)
	private String publicKey;

	@Parameter(displayName = "JSON值", type = ParameterType.MULTILINE_TEXT, order = 6)
	private String json = "{\"key\":\"value\"}";

	@Parameter(displayName = "过期秒数", description = "留空表示不会过期", type = ParameterType.TEXT, order = 7)
	private String expire = "1800";

	@Parameter(displayName = "Jwt值", type = ParameterType.TEXT, order = 8, visible = false)
	private String jwt;

	private ObjectMapper objectMapper = new ObjectMapper();

	@OnParameterChange("algorithm")
	public void onAlgorithmChange(Tool tool) {
		if (algorithm.startsWith("H")) {
			tool.getParameter("key").setVisible(true);
			tool.getParameter("privateKey").setVisible(false);
			tool.getParameter("publicKey").setVisible(false);
		} else {
			tool.getParameter("key").setVisible(false);
			if ("Generate".equals(operation)) {
				tool.getParameter("privateKey").setVisible(true);
				tool.getParameter("publicKey").setVisible(false);
			} else {
				tool.getParameter("privateKey").setVisible(false);
				tool.getParameter("publicKey").setVisible(true);
			}
		}
		if ("Generate".equals(operation)) {
			tool.getParameter("json").setVisible(true);
			tool.getParameter("expire").setVisible(true);
			tool.getParameter("jwt").setVisible(false);
			tool.getAction("generateJwt").setVisible(true);
			tool.getAction("parseJwt").setVisible(false);
		} else {
			tool.getParameter("json").setVisible(false);
			tool.getParameter("expire").setVisible(false);
			tool.getParameter("jwt").setVisible(true);
			tool.getAction("generateJwt").setVisible(false);
			tool.getAction("parseJwt").setVisible(true);
		}
	}

	@OnParameterChange("operation")
	public void onOperationChange(Tool tool) {
		onAlgorithmChange(tool);
	}

	@Action(displayName = "生成JWT", order = 1)
	@SuppressWarnings("unchecked")
	public void generateJwt(Tool tool) throws Exception {
		tool.clearMessages();
		Key skey = null;
		if (algorithm.startsWith("H")) {
			if (CommonUtils.isBlank(key)) {
				tool.errorMessage("必须设置密钥");
				return;
			}
			skey = getHmacKeyFromString(algorithm, key);
		} else {
			if (CommonUtils.isBlank(privateKey)) {
				tool.errorMessage("必须设置私钥");
				return;
			}
			if (algorithm.startsWith("R")) {
				skey = getPrivateKeyFromString("RSA", privateKey);
			} else {
				skey = getPrivateKeyFromString("EC", privateKey);
			}
		}
		SecureDigestAlgorithm<Key, Key> sda = (SecureDigestAlgorithm<Key, Key>) Jwts.SIG.get().get(algorithm);
		if (CommonUtils.isBlank(json)) {
			tool.errorMessage("必须录入JSON值");
			return;
		}

		Map<String, Object> claims = objectMapper.readValue(json, Map.class);
		JwtBuilder builder = Jwts.builder().claims(claims);
		if (CommonUtils.isNotBlank(expire)) {
			Date expiration = new Date(System.currentTimeMillis() + Integer.parseInt(expire) * 1000);
			builder.expiration(expiration);
		}
		String jwt = builder.signWith(skey, sda).compact();
		tool.infoMessage(jwt);
	}

	@Action(displayName = "解析JWT", order = 2, visible = false)
	public void parseJwt(Tool tool) throws Exception {
		tool.clearMessages();
		if (algorithm.startsWith("H")) {
			if (CommonUtils.isBlank(key)) {
				tool.errorMessage("必须设置密钥");
				return;
			}
			SecretKey skey = getHmacKeyFromString(algorithm, key);
			Jws<Claims> jws = Jwts.parser().verifyWith(skey).build().parseSignedClaims(jwt);
			String claimsJson = objectMapper.writeValueAsString(jws.getPayload());
			tool.infoMessage(claimsJson);
		} else {
			if (CommonUtils.isBlank(privateKey)) {
				tool.errorMessage("必须设置公钥");
				return;
			}
			PublicKey pkey;
			if (algorithm.startsWith("R")) {
				pkey = getPublicKeyFromString("RSA", publicKey);
			} else {
				pkey = getPublicKeyFromString("EC", publicKey);
			}
			Jws<Claims> jws = Jwts.parser().verifyWith(pkey).build().parseSignedClaims(jwt);
			String claimsJson = objectMapper.writeValueAsString(jws.getPayload());
			tool.infoMessage(claimsJson);
		}
	}

	@Action(displayName = "生成密钥", order = 3)
	public void generateKey(Tool tool) throws Exception {
		tool.clearMessages();
		if (algorithm.startsWith("H")) {
			String generatedKey = generateHmacKey();
			this.key = generatedKey;
			tool.infoMessage("密钥: ");
			tool.infoMessage(generatedKey);
		} else {
			String keyAlgorithm = "RSA";
			int size = 2048;
			if (algorithm.startsWith("E")) {
				keyAlgorithm = "EC";
			}
			if ("RS384".equals(algorithm)) {
				size = 3072;
			} else if ("RS512".equals(algorithm)) {
				size = 4096;
			} else if ("ES256".equals(algorithm)) {
				size = 256;
			} else if ("ES384".equals(algorithm)) {
				size = 384;
			} else if ("ES512".equals(algorithm)) {
				size = 521;
			}
			String[] generatedKeyPair = generateKeyPair(keyAlgorithm, size);
			this.privateKey = generatedKeyPair[0];
			this.publicKey = generatedKeyPair[1];
			tool.infoMessage("私钥: ");
			tool.infoMessage(generatedKeyPair[0]);
			tool.infoMessage("公钥: ");
			tool.infoMessage(generatedKeyPair[1]);
		}
	}

	private SecretKey getHmacKeyFromString(String algorithm, String key) {
		SecretKey skey = null;
		byte[] keyBytes = Base64.getDecoder().decode(key);
		if ("HS256".equals(algorithm)) {
			skey = new SecretKeySpec(keyBytes, "HmacSHA256");
		} else if ("HS384".equals(algorithm)) {
			skey = new SecretKeySpec(keyBytes, "HmacSHA384");
		} else if ("HS512".equals(algorithm)) {
			skey = new SecretKeySpec(keyBytes, "HmacSHA512");
		}
		return skey;
	}

	private PrivateKey getPrivateKeyFromString(String keyAlgorithm, String privateKey) {
		try {
			KeyFactory keyFactory = KeyFactory.getInstance(keyAlgorithm);
			byte[] encodedKey = Base64.getDecoder().decode(privateKey);
			PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(encodedKey);
			return keyFactory.generatePrivate(keySpec);
		} catch (Exception e) {
			throw new IllegalArgumentException("Invalid private key string: " + privateKey, e);
		}
	}

	private PublicKey getPublicKeyFromString(String keyAlgorithm, String publicKey) {
		try {
			KeyFactory keyFactory = KeyFactory.getInstance(keyAlgorithm);
			byte[] encodedKey = Base64.getDecoder().decode(publicKey);
			X509EncodedKeySpec keySpec = new X509EncodedKeySpec(encodedKey);
			return keyFactory.generatePublic(keySpec);
		} catch (Exception e) {
			throw new IllegalArgumentException("Invalid public key string: " + publicKey, e);
		}
	}

	private String generateHmacKey() {
		MacAlgorithm ma = null;
		if ("HS256".equals(algorithm)) {
			ma = Jwts.SIG.HS256;
		} else if ("HS384".equals(algorithm)) {
			ma = Jwts.SIG.HS384;
		} else if ("HS512".equals(algorithm)) {
			ma = Jwts.SIG.HS512;
		}
		byte[] encoded = ma.key().build().getEncoded();
		String base64 = Base64.getEncoder().encodeToString(encoded);
		return base64;
	}

	private String[] generateKeyPair(String keyAlgorithm, int size) throws NoSuchAlgorithmException {
		KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(keyAlgorithm);
		keyPairGenerator.initialize(size);
		KeyPair keyPair = keyPairGenerator.generateKeyPair();
		PrivateKey privateKey = keyPair.getPrivate();
		byte[] privateKeyEncoded = privateKey.getEncoded();
		String privateKeyBase64 = Base64.getEncoder().encodeToString(privateKeyEncoded);
		PublicKey publicKey = keyPair.getPublic();
		byte[] publicKeyEncoded = publicKey.getEncoded();
		String publicKeyBase64 = Base64.getEncoder().encodeToString(publicKeyEncoded);
		return new String[] { privateKeyBase64, publicKeyBase64 };
	}

}
